import numpy as np
import pandas as pd
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple
import numba
from scipy.optimize import linear_sum_assignment
from collections import deque
import time
from main import HighPrecisionThreatAssessment, RadarDetection, PlatParaInfo
import math


@dataclass
class TrackInfo:
    """敌方航迹信息"""
    fLat: float  # 纬度
    fLon: float  # 经度  
    fHeight: float  # 高度
    fSpeed: float  # 速度
    fHeading: float  # 航向
    trackID: int  # 航迹ID
    timestamp: float  # 时间戳
    target_type: int = 0  # 目标类型
    vulnerability: float = 0.5  # 易损性

class TargetAllocationSystem:
    """高精度实时目标分配系统"""
    
    def __init__(self, threat_assessor):
        self.threat_assessor = threat_assessor
        
        # 分配参数
        self.allocation_params = self._initialize_allocation_params()
        
        # 缓存和状态
        self.platform_capabilities = {}
        self.weapon_inventory = {}
        self.allocation_history = {}
        
        # 性能优化
        self.max_candidates = 50
        self.cache_size = 100
        
    def _initialize_allocation_params(self):
        """初始化分配参数"""
        return {
            # 权重参数
            'alpha': 0.4,  # 威胁权重
            'beta': 0.3,   # 拦截概率权重  
            'gamma': 0.2,  # 效能权重
            'delta': 0.1,  # 成本权重
            
            # 拦截参数
            'R_min': 1000,    # 最小拦截距离
            'R_max': 50000,   # 最大拦截距离
            'theta_max': np.radians(60),  # 最大进入角
            
            # 时间参数
            'decision_timeout': 2.0,  # 决策超时时间
            'reallocation_threshold': 0.3
        }
    
    def calculate_intercept_feasibility(self, platform: PlatParaInfo, target: TrackInfo):
        """计算拦截可行性（修改后的版本）"""
        # 计算相对位置
        radar_detection = RadarDetection(
            fLat=target.fLat, fLon=target.fLon, fHeight=target.fHeight,
            timestamp=target.timestamp, ID=target.trackID
        )
    
        enu = self.threat_assessor.calculate_relative_position(platform, radar_detection)
        R = np.linalg.norm(enu)
    
        # 距离可行性
        if R < self.allocation_params['R_min'] or R > self.allocation_params['R_max']:
            return 0.0, float('inf')
    
        # 使用基本数据类型调用Numba函数
        V_rel = self._calculate_relative_velocity(
            platform_speed=platform.fSpeed,
            platform_yaw=platform.fYaw,
            platform_pitch=platform.fPitch,
            target_speed=target.fSpeed,
            target_heading=target.fHeading,
            target_pitch=getattr(target, 'fPitch', 0.0)  # 默认俯仰角为0
        )
    
        # 将元组转换为numpy数组
        V_rel_array = np.array(V_rel)
        V_rel_norm = np.linalg.norm(V_rel_array)
    
        # 速度可行性
        if V_rel_norm < 1e-6:
            return 0.0, float('inf')
    
        # 计算拦截时间
        intercept_time = R / V_rel_norm
    
        # 进入角评估
        aspect_angle = self._calculate_aspect_angle(enu, V_rel_array)
        aspect_factor = max(0, 1 - abs(aspect_angle) / self.allocation_params['theta_max'])
    
        # 拦截概率
        P_intercept = self._calculate_intercept_probability(R, intercept_time, aspect_factor)
    
        return P_intercept, intercept_time

    @staticmethod
    @numba.jit(nopython=True, fastmath=True, cache=True)
    def _calculate_relative_velocity(platform_speed, platform_yaw, platform_pitch,
                               target_speed, target_heading, target_pitch=0.0):
        """计算相对速度（使用基本数据类型，Numba兼容版本）"""
        deg_to_rad = 0.017453292519943295  # π/180
    
        # 转换为弧度
        platform_heading_rad = platform_yaw * deg_to_rad
        platform_pitch_rad = platform_pitch * deg_to_rad
        target_heading_rad = target_heading * deg_to_rad
        target_pitch_rad = target_pitch * deg_to_rad
    
        # 计算三角函数
        sin_platform_h = math.sin(platform_heading_rad)
        cos_platform_h = math.cos(platform_heading_rad)
        sin_platform_p = math.sin(platform_pitch_rad)
        cos_platform_p = math.cos(platform_pitch_rad)   
    

        # 计算三角函数
        sin_target_h = math.sin(target_heading_rad)
        cos_target_h = math.cos(target_heading_rad)
        sin_target_p = math.sin(target_pitch_rad)
        cos_target_p = math.cos(target_pitch_rad)
    
        # 计算速度分量
        Vx_platform = platform_speed * cos_platform_p * sin_platform_h
        Vy_platform = platform_speed * cos_platform_p * cos_platform_h
        Vz_platform = platform_speed * sin_platform_p
        
        Vx_target = target_speed * cos_target_p * sin_target_h
        Vy_target = target_speed * cos_target_p * cos_target_h
        Vz_target = target_speed * sin_target_p
    
        # 相对速度
        Vx_rel = Vx_target - Vx_platform
        Vy_rel = Vy_target - Vy_platform
        Vz_rel = Vz_target - Vz_platform
        
        return Vx_rel, Vy_rel, Vz_rel
    
    def _calculate_aspect_angle(self, rel_position, rel_velocity):
        """计算进入角"""
        if np.linalg.norm(rel_position) < 1e-6 or np.linalg.norm(rel_velocity) < 1e-6:
            return 0.0
        
        # 归一化向量
        pos_norm = rel_position / np.linalg.norm(rel_position)
        vel_norm = rel_velocity / np.linalg.norm(rel_velocity)
        
        # 计算夹角
        cos_angle = np.dot(pos_norm, vel_norm)
        cos_angle = np.clip(cos_angle, -1.0, 1.0)
        
        return np.arccos(cos_angle)
    
    def _calculate_intercept_probability(self, distance, intercept_time, aspect_factor):
        """计算拦截概率"""
        # 距离因素
        distance_factor = 1.0 / (1.0 + np.exp(0.0001 * (distance - 20000)))
        
        # 时间因素
        time_factor = 1.0 / (1.0 + np.exp(0.1 * (intercept_time - 60)))
        
        # 综合拦截概率
        P_intercept = distance_factor * time_factor * aspect_factor
        
        return np.clip(P_intercept, 0.0, 1.0)
    
    def calculate_effectiveness_coefficient(self, platform_id, target_id, 
                                         platform_info, target_info, threat_value):
        """计算效能系数"""
        params = self.allocation_params
        
        # 拦截概率
        P_intercept, intercept_time = self.calculate_intercept_feasibility(platform_info, target_info)
        
        # 武器目标匹配度
        E_match = self._calculate_weapon_target_match(platform_id, target_info)
        
        # 成本系数（基于距离和时间）
        C_cost = self._calculate_engagement_cost(platform_info, target_info, intercept_time)
        
        # 综合效能系数
        w_ij = (params['alpha'] * threat_value + 
                params['beta'] * P_intercept + 
                params['gamma'] * E_match - 
                params['delta'] * C_cost)
        
        return max(0, w_ij), P_intercept, intercept_time
    
    def _calculate_weapon_target_match(self, platform_id, target_info):
        """计算武器目标匹配度"""
        # 获取平台能力
        platform_capability = self.platform_capabilities.get(platform_id, {})
        weapon_type = platform_capability.get('weapon_type', 0)
        
        # 目标易损性
        target_vulnerability = target_info.vulnerability
        
        # 匹配度计算
        match_score = 1.0 / (1.0 + np.exp(-2.0 * (weapon_type - target_vulnerability)))
        
        return match_score
    
    def _calculate_engagement_cost(self, platform_info, target_info, intercept_time):
        """计算交战成本"""
        # 距离成本
        radar_detection = RadarDetection(
            fLat=target_info.fLat, fLon=target_info.fLon, 
            fHeight=target_info.fHeight, timestamp=target_info.timestamp, 
            ID=target_info.trackID
        )
        enu = self.threat_assessor.calculate_relative_position(platform_info, radar_detection)
        distance_cost = np.linalg.norm(enu) / 10000.0  # 归一化
        
        # 时间成本
        time_cost = intercept_time / 120.0  # 归一化
        
        # 风险成本（基于威胁）
        threats = self.threat_assessor.update_threat_assessment(platform_info, [radar_detection])
        threat_cost = threats.get(target_info.trackID, {}).get('total_threat', 0.5)
        
        total_cost = 0.4 * distance_cost + 0.3 * time_cost + 0.3 * threat_cost
        
        return total_cost
    
    def dynamic_priority_calculation(self, target_info, threat_value, current_time):
        """动态优先级计算"""
        # 威胁紧迫性
        time_criticality = self._calculate_time_criticality(target_info, current_time)
        
        # 战略价值
        strategic_value = self._calculate_strategic_value(target_info)
        
        # 机动性威胁
        maneuver_threat = self._calculate_maneuver_threat(target_info)
        
        # 综合优先级
        priority = (0.5 * threat_value + 
                   0.3 * time_criticality + 
                   0.15 * strategic_value + 
                   0.05 * maneuver_threat)
        
        return priority
    
    def _calculate_time_criticality(self, target_info, current_time):
        """计算时间紧迫性"""
        # 基于目标速度和距离的紧急程度
        time_to_engage = 60.0  # 假设的临界时间
        time_factor = 1.0 / (1.0 + np.exp(-0.1 * (time_to_engage - 30)))
        
        return time_factor
    
    def _calculate_strategic_value(self, target_info):
        """计算战略价值"""
        # 基于目标类型的重要性
        target_type_weights = {
            0: 0.5,   # 普通目标
            1: 0.8,   # 重要目标
            2: 1.0    # 高价值目标
        }
        
        return target_type_weights.get(target_info.target_type, 0.5)
    
    def _calculate_maneuver_threat(self, target_info):
        """计算机动性威胁"""
        # 基于速度和机动能力的威胁评估
        speed_threat = min(1.0, target_info.fSpeed / 300.0)  # 归一化
        
        return speed_threat
    
    def hierarchical_allocation_optimization(self, platforms: List[PlatParaInfo], 
                                          tracks: List[TrackInfo]):
        """分层优化分配"""
        start_time = time.time()
        
        # 第一层：快速筛选候选对
        candidate_pairs = self._fast_candidate_screening(platforms, tracks)
        
        if not candidate_pairs:
            return {}
        
        # 第二层：精确效能评估
        score_matrix, feasibility_info = self._precise_scoring(candidate_pairs, platforms, tracks)
        
        # 第三层：优化分配
        allocation = self._optimization_allocation(score_matrix, candidate_pairs, 
                                                 feasibility_info, platforms, tracks)
        
        decision_time = time.time() - start_time
        print(f"分配决策时间: {decision_time:.3f}秒")
        
        return allocation
    
    def _fast_candidate_screening(self, platforms, tracks):
        """快速候选对筛选"""
        candidate_pairs = []
        
        for i, platform in enumerate(platforms):
            for j, track in enumerate(tracks):
                # 基本可行性检查
                radar_detection = RadarDetection(
                    fLat=track.fLat, fLon=track.fLon, fHeight=track.fHeight,
                    timestamp=track.timestamp, ID=track.trackID
                )
                
                enu = self.threat_assessor.calculate_relative_position(platform, radar_detection)
                R = np.linalg.norm(enu)
                
                # 距离筛选
                if R < self.allocation_params['R_min'] or R > self.allocation_params['R_max']:
                    continue
                
                # 拦截概率初步评估
                P_intercept, _ = self.calculate_intercept_feasibility(platform, track)
                if P_intercept < 0.1:  # 最低拦截概率阈值
                    continue
                
                candidate_pairs.append((i, j))
                
                # 限制候选对数量
                if len(candidate_pairs) >= self.max_candidates:
                    return candidate_pairs
        
        return candidate_pairs
    
    def _precise_scoring(self, candidate_pairs, platforms, tracks):
        """精确效能评分"""
        n_platforms = len(platforms)
        n_targets = len(tracks)
        
        # 初始化评分矩阵
        score_matrix = np.full((n_platforms, n_targets), -np.inf)
        feasibility_info = {}
        
        # 计算威胁值
        threat_values = {}
        for j, track in enumerate(tracks):
            # 使用最近平台作为参考计算威胁
            if platforms:
                radar_detection = RadarDetection(
                    fLat=track.fLat, fLon=track.fLon, fHeight=track.fHeight,
                    timestamp=track.timestamp, ID=track.trackID
                )
                threats = self.threat_assessor.update_threat_assessment(platforms[0], [radar_detection])
                threat_values[j] = threats.get(track.trackID, {}).get('total_threat', 0.5)
        
        for i, j in candidate_pairs:
            platform = platforms[i]
            track = tracks[j]
            
            threat_value = threat_values.get(j, 0.5)
            
            # 计算效能系数
            w_ij, P_intercept, intercept_time = self.calculate_effectiveness_coefficient(
                platform.ID, track.trackID, platform, track, threat_value
            )
            
            # 计算动态优先级
            priority = self.dynamic_priority_calculation(track, threat_value, time.time())
            
            # 综合评分
            final_score = w_ij * priority
            
            score_matrix[i, j] = final_score
            feasibility_info[(i, j)] = {
                'intercept_probability': P_intercept,
                'intercept_time': intercept_time,
                'threat_value': threat_value
            }
        
        return score_matrix, feasibility_info
    
    def _optimization_allocation(self, score_matrix, candidate_pairs, feasibility_info, 
                               platforms, tracks):
        """优化分配求解"""
        n_platforms, n_targets = score_matrix.shape
        
        # 资源约束：每个平台最多分配的目标数
        platform_capacity = self._get_platform_capacity(platforms)
        
        # 使用匈牙利算法进行分配
        try:
            # 处理负无穷大值
            valid_matrix = np.where(score_matrix == -np.inf, -1e9, score_matrix)
            
            # 行规约
            row_reduced = valid_matrix - np.max(valid_matrix, axis=1, keepdims=True)
            
            # 列规约
            cost_matrix = -row_reduced  # 转换为最小化问题
            
            # 使用匈牙利算法
            row_ind, col_ind = linear_sum_assignment(cost_matrix)
            
            allocation = {}
            for i, j in zip(row_ind, col_ind):
                if score_matrix[i, j] > 0:  # 只接受正分数的分配
                    platform_id = platforms[i].ID
                    target_id = tracks[j].trackID
                    
                    allocation[platform_id] = {
                        'target_id': target_id,
                        'score': score_matrix[i, j],
                        'intercept_probability': feasibility_info.get((i, j), {}).get('intercept_probability', 0),
                        'intercept_time': feasibility_info.get((i, j), {}).get('intercept_time', 0),
                        'threat_value': feasibility_info.get((i, j), {}).get('threat_value', 0)
                    }
            
            return allocation
            
        except Exception as e:
            print(f"优化算法错误: {e}")
            return self._heuristic_allocation(score_matrix, platforms, tracks, feasibility_info)
    
    def _heuristic_allocation(self, score_matrix, platforms, tracks, feasibility_info):
        """启发式分配（备用方法）"""
        allocation = {}
        n_platforms, n_targets = score_matrix.shape
        
        # 按评分排序
        assignments = []
        for i in range(n_platforms):
            for j in range(n_targets):
                if score_matrix[i, j] > 0:
                    assignments.append((i, j, score_matrix[i, j]))
        
        # 按评分降序排序
        assignments.sort(key=lambda x: x[2], reverse=True)
        
        allocated_platforms = set()
        allocated_targets = set()
        
        for i, j, score in assignments:
            if i not in allocated_platforms and j not in allocated_targets:
                platform_id = platforms[i].ID
                target_id = tracks[j].trackID
                
                allocation[platform_id] = {
                    'target_id': target_id,
                    'score': score,
                    'intercept_probability': feasibility_info.get((i, j), {}).get('intercept_probability', 0),
                    'intercept_time': feasibility_info.get((i, j), {}).get('intercept_time', 0),
                    'threat_value': feasibility_info.get((i, j), {}).get('threat_value', 0)
                }
                
                allocated_platforms.add(i)
                allocated_targets.add(j)
        
        return allocation
    
    def _get_platform_capacity(self, platforms):
        """获取平台容量"""
        capacities = {}
        for platform in platforms:
            # 基于平台类型和能力确定容量
            capacities[platform.ID] = 2  # 默认每个平台分配2个目标
        return capacities
    
    def time_sensitive_decision(self, platforms, tracks, time_constraint=2.0):
        """时间敏感决策"""
        start_time = time.time()
        
        # 快速威胁评估
        quick_threats = self._quick_threat_assessment(platforms[0] if platforms else None, tracks)
        
        # 紧急目标识别
        urgent_targets = self._identify_urgent_targets(tracks, quick_threats, time_constraint)
        
        # 快速分配
        if urgent_targets:
            quick_allocation = self._quick_emergency_allocation(platforms, urgent_targets)
            
            decision_time = time.time() - start_time
            if decision_time < time_constraint:
                return quick_allocation
        
        # 如果时间允许，进行完整优化
        return self.hierarchical_allocation_optimization(platforms, tracks)
    
    def _quick_threat_assessment(self, reference_platform, tracks):
        """快速威胁评估"""
        threats = {}
        
        if reference_platform is None:
            return {track.trackID: 0.5 for track in tracks}
        
        for track in tracks:
            radar_detection = RadarDetection(
                fLat=track.fLat, fLon=track.fLon, fHeight=track.fHeight,
                timestamp=track.timestamp, ID=track.trackID
            )
            
            enu = self.threat_assessor.calculate_relative_position(reference_platform, radar_detection)
            R = np.linalg.norm(enu)
            
            # 简化的威胁计算
            threat_value = 1.0 / (1.0 + 0.0001 * R)  # 基于距离的简化威胁
            threats[track.trackID] = threat_value
        
        return threats
    
    def _identify_urgent_targets(self, tracks, threats, time_constraint):
        """识别紧急目标"""
        urgent_targets = []
        
        for track in tracks:
            threat_value = threats.get(track.trackID, 0)
            
            # 基于速度和威胁的紧急程度
            urgency_score = threat_value * (track.fSpeed / 200.0)
            
            if urgency_score > 0.7:  # 紧急阈值
                urgent_targets.append(track)
        
        return urgent_targets
    
    def _quick_emergency_allocation(self, platforms, urgent_targets):
        """紧急情况快速分配"""
        allocation = {}
        
        # 简单的一对一分配
        for i, platform in enumerate(platforms):
            if i < len(urgent_targets):
                target = urgent_targets[i]
                
                allocation[platform.ID] = {
                    'target_id': target.trackID,
                    'score': 1.0,  # 紧急分配满分
                    'intercept_probability': 0.8,  # 假设值
                    'intercept_time': 60.0,  # 假设值
                    'threat_value': 0.8,  # 假设值
                    'emergency': True
                }
        
        return allocation
    
    def update_platform_capabilities(self, platform_id, capabilities):
        """更新平台能力信息"""
        self.platform_capabilities[platform_id] = capabilities
    
    def reallocation_trigger(self, current_allocation, new_threats, platforms, tracks):
        """重分配触发判断"""
        if not current_allocation:
            return True
        
        # 威胁变化检测
        max_threat_change = 0.0
        for platform_id, assignment in current_allocation.items():
            target_id = assignment['target_id']
            old_threat = assignment.get('threat_value', 0)
            new_threat = new_threats.get(target_id, 0)
            
            threat_change = abs(new_threat - old_threat)
            max_threat_change = max(max_threat_change, threat_change)
        
        # 拦截概率变化检测
        min_intercept_prob = 1.0
        for platform in platforms:
            for track in tracks:
                P_intercept, _ = self.calculate_intercept_feasibility(platform, track)
                min_intercept_prob = min(min_intercept_prob, P_intercept)
        
        # 触发条件
        trigger = (max_threat_change > self.allocation_params['reallocation_threshold'] or 
                  min_intercept_prob < 0.3)
        
        return trigger

# 使用示例
def main():
    # 创建威胁评估系统
    threat_assessor = HighPrecisionThreatAssessment()
    
    # 创建目标分配系统
    allocation_system = TargetAllocationSystem(threat_assessor)
    
    # 模拟数据
    platforms = [
        PlatParaInfo(fLat=39.9, fLon=116.4, fHeight=1000, fSpeed=200,
                    fRoll=0, fPitch=0, fYaw=0, ID=1, timestamp=time.time()),
        PlatParaInfo(fLat=39.91, fLon=116.41, fHeight=1100, fSpeed=180,
                    fRoll=0, fPitch=0, fYaw=10, ID=2, timestamp=time.time())
    ]
    
    tracks = [
        TrackInfo(fLat=39.92, fLon=116.42, fHeight=1200, fSpeed=250,
                 fHeading=45, trackID=1001, timestamp=time.time(), target_type=1),
        TrackInfo(fLat=39.88, fLon=116.38, fHeight=900, fSpeed=300,
                 fHeading=120, trackID=1002, timestamp=time.time(), target_type=2),
        TrackInfo(fLat=39.93, fLon=116.43, fHeight=1500, fSpeed=200,
                 fHeading=300, trackID=1003, timestamp=time.time(), target_type=0)
    ]
    
    # 设置平台能力
    allocation_system.update_platform_capabilities(1, {'weapon_type': 0.8, 'capacity': 2})
    allocation_system.update_platform_capabilities(2, {'weapon_type': 0.6, 'capacity': 1})
    
    # 进行目标分配
    allocation = allocation_system.hierarchical_allocation_optimization(platforms, tracks)
    
    # 输出结果
    print("目标分配结果:")
    for platform_id, assignment in allocation.items():
        print(f"平台 {platform_id} -> 目标 {assignment['target_id']}")
        print(f"  评分: {assignment['score']:.3f}")
        print(f"  拦截概率: {assignment['intercept_probability']:.3f}")
        print(f"  拦截时间: {assignment['intercept_time']:.1f}秒")
        print(f"  威胁值: {assignment['threat_value']:.3f}")

if __name__ == "__main__":
    main()